/*******************************************************************************
 *  SarsaLAMBDAAgent.cpp
 *  
 *  Created by Adam White, created on April 18 2007.
 *  Copyright (c) 2007 UAlberta. All rights reserved.
 *
 *  Description.
 *
 *		A Sarsa LAMBDA control Agent with Tile coding Function Approximation. 
 *		This agent has been tuned for the Acrobot task. It is based on 
 *		Sutton's code, found here:
 *
 *       http://www.cs.ualberta.ca/~sutton/MountainCar/MountainCar2.cpp
 *
 *      The code follows the psuedo-code for linear, gradient-descent 
 *		Sarsa(LAMBDA) given in Figure 8.8 of the book "Reinforcement Learning: 
 *		An Introduction", by Sutton and Barto. One difference is that we use 
 *		the implementation trick mentioned on page 189 to only keep track of 
 *		the traces that are larger than "min-trace". 
 *					
 ********************************************************************************/

#include "SarsaLambdaAgent.h"
#include <string.h>
#include "ParamReader.h"



//-----------------variables-----------------------
/*
double const TILE_WIDTH_0 = M_PI/3.0;               // the tile width for position
double const TILE_WIDTH_1 = (0.14 / 8);              // the tile width for velocity
double const TILE_WIDTH_2 = 8.0*M_PI/6.0;
double const TILE_WIDTH_3 = 3.0*M_PI;
*/
//ali version:

/*double const TILE_WIDTH_0 = M_PI/4.0;               // the tile width for position
double const TILE_WIDTH_1 = M_PI/4.0;              // the tile width for velocity
double const TILE_WIDTH_2 = M_PI;
double const TILE_WIDTH_3 = 9*M_PI/4.0;
*/
double* tileWidth; 

// Global RL variables:
int NUM_ACTIONS;                         //number of available actions   
int NUM_OBSERVATIONS;                     //dimension of observation

int NUM_TILINGS; 						  //number of tiling grids	
double EPSILON;                    // probability of random action
double ALPHA;                      // step size parameter
double LAMBDA;                     // trace-decay parameters

double* QSA;								// action values
double* weights;							// feature weights
double* traces;                             // eligibility traces
int** activeFeatures;           // sets of features, one set per action

// Standard RL parameters:
double RLgamma=1.0;                // discount-rate parameters

action_t oldAction;						//action selected on previous time step
action_t newAction;						//action selected on current time step


taskspec_t *ts;					//task spec


//sutton 
int nonzero_traces[MAX_NONZERO_TRACES];
int* nonzero_traces_inverse;

double minimum_trace = 0.01;
int num_nonzero_traces = 0;


//--------------------------------------------------




// RLGLUE INTERFACE FUNCTIONS ----------------------------------------------
void agent_init(const char* _ts)
{

	EPSILON = ParamReader::AGNT_EPSILON;                    // probability of random action
	ALPHA = ParamReader::AGNT_ALPHA;                      // step size parameter
	LAMBDA = ParamReader::AGNT_LAMBDA;                    // trace-decay parameters	
	NUM_TILINGS = ParamReader::AGNT_NUMBER_OF_TILINGS;  


	ts=(taskspec_t*)malloc(sizeof(taskspec_t));

	allocateRLStruct(&oldAction, 1, 0, 0); 
	allocateRLStruct(&newAction, 1, 0, 0); 

	int decode_result = decode_taskspec( ts, _ts);
	if(decode_result!=0){
		printf("Could not decode task spec, code: %d for task spec: %s\n",decode_result,_ts);
		exit(1);
	}

	//initialize some variables 
	NUM_ACTIONS = ts->int_actions[0].max - ts->int_actions[0].min + 1; 
	NUM_OBSERVATIONS = ts->num_double_observations; 

	
	//create containers
	QSA = new double [NUM_ACTIONS]; 
	weights = new double[MEMORY_SIZE]; 
	traces = new double [MEMORY_SIZE]; 
	nonzero_traces_inverse = new int[MEMORY_SIZE]; 
	activeFeatures = Allocate2DArray<int>(NUM_ACTIONS, NUM_TILINGS); 
	tileWidth = new double[NUM_OBSERVATIONS]; 

	for(int i=0; i< NUM_OBSERVATIONS; i++)
		tileWidth[i] = (ts->double_observations[i].max - ts->double_observations[i].min)/8.0; 


	memset(weights, 0, MEMORY_SIZE*sizeof(double)); //clear weights
	memset(traces, 0, MEMORY_SIZE*sizeof(double));	//clear all traces

}

void agent_cleanup()
{
	clearRLStruct(&oldAction); 
	clearRLStruct(&newAction); 

	delete[] QSA; 
	delete[] weights; 
	Free2DArray(activeFeatures); 
}





const action_t* agent_start(const observation_t* o)
{
    DecayTraces(0.0);                                           
    computeActiveFeatures(o);                                                     
    computeActionValues(); 
	
	selectEpsilonGreedyAction(o,oldAction);
                                     
	return &oldAction;
}

const action_t* agent_step(double r, const observation_t* o)
{
	updateTraces();
	
	computeActionValues(oldAction);    //compute value of Q[oldObservation][oldAction]
	double delta = r - QSA[oldAction.intArray[0]];
	
	computeActiveFeatures(o);                                              
	computeActionValues();			  //new action values based on new observation											
	
	selectEpsilonGreedyAction(o,newAction);

	delta += RLgamma * QSA[newAction.intArray[0]];
	
    updateWeights(delta);

	oldAction.intArray[0] = newAction.intArray[0]; 


	return &newAction;
}

void agent_end(double r)
{
	updateTraces();
		
	computeActionValues(oldAction);
	double delta = r - QSA[oldAction.intArray[0]];
			
	updateWeights(delta);
	
	EPSILON *= 0.99;
}

 
void agent_freeze()
{
	ALPHA = 0.0;
	LAMBDA = 0.0;
} 

const char* agent_message(const char*){
return "There are currently no messages implemented for agent_message";
}
 
// HELPER FUNCTIONS ----------------------------------------------
 
void selectEpsilonGreedyAction(const observation_t* o, action_t& action)
{
//select an action according to EPSILON greedy policy 

	if(drand48() <= EPSILON) 
	{
		action.intArray[0] = (int)(drand48() * (NUM_ACTIONS));
	}
	else
		action.intArray[0] = argmax(QSA); 

	if (action.intArray[0] >= NUM_ACTIONS)
		action.intArray[0] = NUM_ACTIONS -1 ; 
}

void updateTraces()
{
//At the beggining of every episode, decay traces, clear action traces and replace current trace

	DecayTraces(RLgamma*LAMBDA);                              

	for (int a=0; a<NUM_ACTIONS; a++)                        
		if (a != oldAction.intArray[0])
			for (int j=0; j<NUM_TILINGS; j++) ClearTrace(activeFeatures[a][j]);
        for (int j=0; j<NUM_TILINGS; j++) SetTrace(activeFeatures[oldAction.intArray[0]][j],1.0); // replace traces
}

void updateWeights(double delta)
{
//Update weights with nonzero traces using td-error

	double temp = (ALPHA/NUM_TILINGS)*delta;
	for (int i=0; i<num_nonzero_traces; i++)                
	{ 
		int index = nonzero_traces[i];
		weights[index] += temp * traces[index];
	}  
}

void computeActionValues() 
{
// Compute all the action values from current activeFeatures and weights

	for (int a=0; a<NUM_ACTIONS; a++) 
	{
		QSA[a] = 0;
		for (int j=0; j<NUM_TILINGS; j++) 
			QSA[a] += weights[activeFeatures[a][j]];
	}
}
		 
void computeActionValues(const action_t& a) 
{
// Compute a particular action value from current activeFeatures and weights

	QSA[a.intArray[0]] = 0;
    for (int j=0; j<NUM_TILINGS; j++) 
		QSA[a.intArray[0]] += weights[activeFeatures[a.intArray[0]][j]];
}

void computeActiveFeatures(const observation_t* o)
{
// get set of active features for current observation. One for each action

	double* inputObservations = new double [NUM_OBSERVATIONS];

	for(int i=0; i< NUM_OBSERVATIONS; i++)
		inputObservations[i] = o->doubleArray[i] / tileWidth[i]; 

	for (int a=0; a<NUM_ACTIONS; a++)
        tiles(&activeFeatures[a][0],NUM_TILINGS,MEMORY_SIZE, inputObservations,NUM_OBSERVATIONS,a);

	delete[] inputObservations; 
}

int argmax(double* QSA)
{
// Returns index (action) of largest entry in QSA array, breaking ties randomly

	int best_action = 0;
    double best_value = QSA[0];
    int num_ties = 1;                    // actually the number of ties plus 1
    double value;
	
	for (int a=1; a<NUM_ACTIONS; a++) 
	{
		value = QSA[a];
        if (value >= best_value) 
            if (value > best_value)
			{
				best_value = value;
				best_action = a;
			}
            else 
			{
				num_ties++;
				if (0 == (int)(drand48()*num_ties))
				{
					best_value = value;
					best_action = a;
				}
			}
	}
    return best_action;
}

// ------------------------ Suttons Trace Code -----------------------------------------------------
void SetTrace(int f, double new_trace_value)
// Set the trace for feature f to the given value, which must be positive
  { if (traces[f] >= minimum_trace) traces[f] = new_trace_value;         // trace already exists
    else { while (num_nonzero_traces >= MAX_NONZERO_TRACES) IncreaseMinTrace(); // ensure room for new trace
           traces[f] = new_trace_value;
           nonzero_traces[num_nonzero_traces] = f;
           nonzero_traces_inverse[f] = num_nonzero_traces;
           num_nonzero_traces++;}}

void ClearTrace(int f)       
// Clear any trace for feature f
{ 
	if (!(traces[f]==0.0)) 
        ClearExistentTrace(f,nonzero_traces_inverse[f]); 
}

void ClearExistentTrace(int f, int loc)
// Clear the trace for feature f at location loc in the list of nonzero traces
  { traces[f] = 0.0;
    num_nonzero_traces--;
    nonzero_traces[loc] = nonzero_traces[num_nonzero_traces];
    nonzero_traces_inverse[nonzero_traces[loc]] = loc;}

void DecayTraces(double decay_rate)
// Decays all the (nonzero) traces by decay_rate, removing those below minimum_trace
{ 
	for (int loc=num_nonzero_traces-1; loc>=0; loc--)      // necessary to loop downwards
    { 
		int f = nonzero_traces[loc];
        traces[f] *= decay_rate;
        if (traces[f] < minimum_trace) ClearExistentTrace(f,loc);
	}

}

void IncreaseMinTrace()
// Try to make room for more traces by incrementing minimum_trace by 10%, 
// culling any traces that fall below the new minimum
{
	minimum_trace += 0.1 * minimum_trace;
	for (int loc=num_nonzero_traces-1; loc>=0; loc--)      // necessary to loop downwards
	{ 
		int f = nonzero_traces[loc];
		if (traces[f] < minimum_trace) 
			ClearExistentTrace(f,loc);
	}
}

#ifdef OS_WINDOWS
double drand48()
{
	return  (double) (rand() / (double)(RAND_MAX));
}
#endif


template < typename T >
T **Allocate2DArray( int nRows, int nCols)
{
    T **ppi;
    T *pool;
    T *curPtr;
    //(step 1) allocate memory for array of elements of column

    ppi = new T*[nRows];

    //(step 2) allocate memory for array of elements of each row
    pool = new T [nRows * nCols];

    // Now point the pointers in the right place
    curPtr = pool;
    for( int i = 0; i < nRows; i++)
    {
        *(ppi + i) = curPtr;
         curPtr += nCols;
    }
    return ppi;
}

template < typename T >
void Free2DArray(T** Array)
{
    delete [] *Array;
    delete [] Array;
}